# 基于torch.nn的算子构建有向图
import torch
from torch import nn
from torchvision.models import resnet18, ResNet18_Weights
import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class Vertex(nn.Module):
    def __init__(self, key, op, num = 0, Timing=False, argc=0):
        # Key is the name of the vertex
        # op is the nn operation(op.forward)
        super(Vertex, self).__init__()
        self.id = key
        self.num = num
        self.connectedTo = {}
        self.op = op
        self.Timing = Timing
        self.cost = 0
        self.argc = argc
        
    def decorator_timing(func):
        from functools import wraps
        starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        
        @wraps(func)
        def wrapper(self, *args, **kwargs):
            if self.Timing:
                starter.record()
                out = func(self, *args, **kwargs)
                ender.record()
                torch.cuda.synchronize() # 等待GPU任务完成
                curr_time = starter.elapsed_time(ender) # 从 starter 到 ender 之间用时,单位为毫秒
                print("{0: >10s}\t{1: >10.4f}(ms)".format(self.id, curr_time))
                self.cost = curr_time
            else:
                out = func(self, *args, **kwargs)
            return out
        return wrapper
    
    def addNeighbor(self, nbr, weight = 0):
        '''
        添加邻接的顶点
        :param nbr: 邻接的顶点
        :param weight: 默认为0
        :return:
        '''
        self.connectedTo[nbr] = weight
    
    def getConnections(self):
        '''
        获取邻接的所有顶点
        :return:
        '''
        return self.connectedTo.keys()
    
    def getId(self):
        '''
        获取自身顶点的值
        :return:
        '''
        return self.id
    
    def getWeight(self, nbr):
        '''
        两顶点连接边的权重
        :param nbr:
        :return:
        '''
        return self.connectedTo[nbr]
    
    @decorator_timing
    def forward(self, *x):
        s = 0
        for v in x:
            s += v
        out = self.op(s)
        return out

class Graph:
    '''
    新建一个空图。
    '''
    def __init__(self, Timing=False):
        '''
        初始化邻接表，顶点个数
        '''
        self.vertList = {}
        self.numVertices = 0
        self.time = 0
        self.Timing = Timing
    def addVertex(self, key, op):
        '''
        向图中添加一个顶点实例
        :param key:
        :return: None
        '''
        newVertex = Vertex(key,op, self.numVertices, Timing=self.Timing)
        self.numVertices = self.numVertices + 1
        self.vertList[key] = newVertex
        return newVertex

    def getVertex(self, n):
        '''
        在图中找到名为n的顶点
        :param n:
        :return:
        '''
        if n in self.vertList:
            return self.vertList[n]
        else:
            return None
    def __contains__(self,n):
        '''
        对in进行重载
        :param n:
        :return:
        '''
        return n in self.vertList
    def addEdge(self, f, t, weight =0 , op = None):
        '''
        向图中添加一条带权重cost的有向边，用于连接顶点f和t
        :param f:
        :param t:
        :param weight:
        :return: None
        '''
        if f not in self.vertList:
            nv = self.addVertex(f, op, 0)
        if t not in self.vertList:
            nv = self.addVertex(t, op, 0 )
        self.vertList[f].addNeighbor(self.vertList[t],weight)
        self.vertList[t].argc +=1

    def getVertics(self):
        '''
        以列表形式返回图中所有顶点
        :return: {}.keys()
        '''
        return self.vertList.keys()
    def __iter__(self):
        '''
        对迭代进行重载，方便遍历顶点邻接表
        :return:
        '''
        return iter(self.vertList.values())

class Graph_resnet18(nn.Module):
    def __init__(self, Timing=False):
        super(Graph_resnet18, self).__init__()
        import torchvision
        model = torchvision.models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
        model.eval().to(device)
        self.net = model
        g = Graph(Timing=Timing)
        self.g =g
        g.addVertex('resnet.conv1'               , model.conv1                )
        g.addVertex('resnet.bn1'                 , model.bn1                  )
        g.addVertex('resnet.relu'                , model.relu                 )
        g.addVertex('resnet.maxpool'             , model.maxpool              )

        g.addVertex('resnet.layer1.0.conv1'      , model.layer1[0].conv1      )
        g.addVertex('resnet.layer1.0.bn1'        , model.layer1[0].bn1        )
        g.addVertex('resnet.layer1.0.relu1'      , model.layer1[0].relu       )
        g.addVertex('resnet.layer1.0.conv2'      , model.layer1[0].conv2      )
        g.addVertex('resnet.layer1.0.bn2'        , model.layer1[0].bn2        )
        g.addVertex('resnet.layer1.0.relu2'      , model.layer1[0].relu       )

        g.addVertex('resnet.layer1.1.conv1'      , model.layer1[1].conv1      )
        g.addVertex('resnet.layer1.1.bn1'        , model.layer1[1].bn1        )
        g.addVertex('resnet.layer1.1.relu1'      , model.layer1[1].relu       )
        g.addVertex('resnet.layer1.1.conv2'      , model.layer1[1].conv2      )
        g.addVertex('resnet.layer1.1.bn2'        , model.layer1[1].bn2        )
        g.addVertex('resnet.layer1.1.relu2'      , model.layer1[1].relu       )

        g.addVertex('resnet.layer2.0.downsample' , model.layer2[0].downsample )
        g.addVertex('resnet.layer2.0.conv1'      , model.layer2[0].conv1      )
        g.addVertex('resnet.layer2.0.bn1'        , model.layer2[0].bn1        )
        g.addVertex('resnet.layer2.0.relu1'      , model.layer2[0].relu       )
        g.addVertex('resnet.layer2.0.conv2'      , model.layer2[0].conv2      )
        g.addVertex('resnet.layer2.0.bn2'        , model.layer2[0].bn2        )
        g.addVertex('resnet.layer2.0.relu2'      , model.layer2[0].relu       )

        g.addVertex('resnet.layer2.1.conv1'      , model.layer2[1].conv1      )
        g.addVertex('resnet.layer2.1.bn1'        , model.layer2[1].bn1        )
        g.addVertex('resnet.layer2.1.relu1'      , model.layer2[1].relu       )
        g.addVertex('resnet.layer2.1.conv2'      , model.layer2[1].conv2      )
        g.addVertex('resnet.layer2.1.bn2'        , model.layer2[1].bn2        )
        g.addVertex('resnet.layer2.1.relu2'      , model.layer2[1].relu       )

        g.addVertex('resnet.layer3.0.downsample' , model.layer3[0].downsample )
        g.addVertex('resnet.layer3.0.conv1'      , model.layer3[0].conv1      )
        g.addVertex('resnet.layer3.0.bn1'        , model.layer3[0].bn1        )
        g.addVertex('resnet.layer3.0.relu1'      , model.layer3[0].relu       )
        g.addVertex('resnet.layer3.0.conv2'      , model.layer3[0].conv2      )
        g.addVertex('resnet.layer3.0.bn2'        , model.layer3[0].bn2        )
        g.addVertex('resnet.layer3.0.relu2'      , model.layer3[0].relu       )
    
        g.addVertex('resnet.layer3.1.conv1'      , model.layer3[1].conv1      )
        g.addVertex('resnet.layer3.1.bn1'        , model.layer3[1].bn1        )
        g.addVertex('resnet.layer3.1.relu1'      , model.layer3[1].relu       )
        g.addVertex('resnet.layer3.1.conv2'      , model.layer3[1].conv2      )
        g.addVertex('resnet.layer3.1.bn2'        , model.layer3[1].bn2        )
        g.addVertex('resnet.layer3.1.relu2'      , model.layer3[1].relu       )

        g.addVertex('resnet.layer4.0.downsample' , model.layer4[0].downsample )
        g.addVertex('resnet.layer4.0.conv1'      , model.layer4[0].conv1      )
        g.addVertex('resnet.layer4.0.bn1'        , model.layer4[0].bn1        )
        g.addVertex('resnet.layer4.0.relu1'      , model.layer4[0].relu       )
        g.addVertex('resnet.layer4.0.conv2'      , model.layer4[0].conv2      )
        g.addVertex('resnet.layer4.0.bn2'        , model.layer4[0].bn2        )
        g.addVertex('resnet.layer4.0.relu2'      , model.layer4[0].relu       )
    
        g.addVertex('resnet.layer4.1.conv1'      , model.layer4[1].conv1      )
        g.addVertex('resnet.layer4.1.bn1'        , model.layer4[1].bn1        )
        g.addVertex('resnet.layer4.1.relu1'      , model.layer4[1].relu       )
        g.addVertex('resnet.layer4.1.conv2'      , model.layer4[1].conv2      )
        g.addVertex('resnet.layer4.1.bn2'        , model.layer4[1].bn2        )
        g.addVertex('resnet.layer4.1.relu2'      , model.layer4[1].relu       )

        g.addVertex('resnet.avgpool'             , model.avgpool              )
        g.addVertex('resnet.flatten'             , nn.Flatten()               )
        g.addVertex('resnet.fc'                  , model.fc                   )
        
        
        g.addEdge('resnet.conv1'               ,'resnet.bn1'                 ,0  )
        g.addEdge('resnet.bn1'                 ,'resnet.relu'                ,0  )
        g.addEdge('resnet.relu'                ,'resnet.maxpool'             ,0  )
        g.addEdge('resnet.maxpool'             ,'resnet.layer1.0.conv1'      ,0  )
        g.addEdge('resnet.maxpool'             ,'resnet.layer1.0.relu2'      ,0  )
        g.addEdge('resnet.layer1.0.conv1'      ,'resnet.layer1.0.bn1'        ,0  )
        g.addEdge('resnet.layer1.0.bn1'        ,'resnet.layer1.0.relu1'      ,0  )
        g.addEdge('resnet.layer1.0.relu1'      ,'resnet.layer1.0.conv2'      ,0  )
        g.addEdge('resnet.layer1.0.conv2'      ,'resnet.layer1.0.bn2'        ,0  )
        g.addEdge('resnet.layer1.0.bn2'        ,'resnet.layer1.0.relu2'      ,0  )
        g.addEdge('resnet.layer1.0.relu2'      ,'resnet.layer1.1.conv1'      ,0  )
        g.addEdge('resnet.layer1.0.relu2'      ,'resnet.layer1.1.relu2'      ,0  )
        g.addEdge('resnet.layer1.1.conv1'      ,'resnet.layer1.1.bn1'        ,0  )
        g.addEdge('resnet.layer1.1.bn1'        ,'resnet.layer1.1.relu1'      ,0  )
        g.addEdge('resnet.layer1.1.relu1'      ,'resnet.layer1.1.conv2'      ,0  )
        g.addEdge('resnet.layer1.1.conv2'      ,'resnet.layer1.1.bn2'        ,0  )
        g.addEdge('resnet.layer1.1.bn2'        ,'resnet.layer1.1.relu2'      ,0  )
        g.addEdge('resnet.layer1.1.relu2'      ,'resnet.layer2.0.downsample' ,0  )
        g.addEdge('resnet.layer1.1.relu2'      ,'resnet.layer2.0.conv1'      ,0  )
        g.addEdge('resnet.layer2.0.downsample' ,'resnet.layer2.0.relu2'      ,0  )
        g.addEdge('resnet.layer2.0.conv1'      ,'resnet.layer2.0.bn1'        ,0  )
        g.addEdge('resnet.layer2.0.bn1'        ,'resnet.layer2.0.relu1'      ,0  )
        g.addEdge('resnet.layer2.0.relu1'      ,'resnet.layer2.0.conv2'      ,0  )
        g.addEdge('resnet.layer2.0.conv2'      ,'resnet.layer2.0.bn2'        ,0  )
        g.addEdge('resnet.layer2.0.bn2'        ,'resnet.layer2.0.relu2'      ,0  )
        g.addEdge('resnet.layer2.0.relu2'      ,'resnet.layer2.1.conv1'      ,0  )
        g.addEdge('resnet.layer2.0.relu2'      ,'resnet.layer2.1.relu2'      ,0  )
        g.addEdge('resnet.layer2.1.conv1'      ,'resnet.layer2.1.bn1'        ,0  )
        g.addEdge('resnet.layer2.1.bn1'        ,'resnet.layer2.1.relu1'      ,0  )
        g.addEdge('resnet.layer2.1.relu1'      ,'resnet.layer2.1.conv2'      ,0  )
        g.addEdge('resnet.layer2.1.conv2'      ,'resnet.layer2.1.bn2'        ,0  )
        g.addEdge('resnet.layer2.1.bn2'        ,'resnet.layer2.1.relu2'      ,0  )
        g.addEdge('resnet.layer2.1.relu2'      ,'resnet.layer3.0.downsample' ,0  )
        g.addEdge('resnet.layer2.1.relu2'      ,'resnet.layer3.0.conv1'      ,0  )
        g.addEdge('resnet.layer3.0.downsample' ,'resnet.layer3.0.relu2'      ,0  )
        g.addEdge('resnet.layer3.0.conv1'      ,'resnet.layer3.0.bn1'        ,0  )
        g.addEdge('resnet.layer3.0.bn1'        ,'resnet.layer3.0.relu1'      ,0  )
        g.addEdge('resnet.layer3.0.relu1'      ,'resnet.layer3.0.conv2'      ,0  )
        g.addEdge('resnet.layer3.0.conv2'      ,'resnet.layer3.0.bn2'        ,0  )
        g.addEdge('resnet.layer3.0.bn2'        ,'resnet.layer3.0.relu2'      ,0  )
        g.addEdge('resnet.layer3.0.relu2'      ,'resnet.layer3.1.conv1'      ,0  )
        g.addEdge('resnet.layer3.0.relu2'      ,'resnet.layer3.1.relu2'      ,0  )
        g.addEdge('resnet.layer3.1.conv1'      ,'resnet.layer3.1.bn1'        ,0  )
        g.addEdge('resnet.layer3.1.bn1'        ,'resnet.layer3.1.relu1'      ,0  )
        g.addEdge('resnet.layer3.1.relu1'      ,'resnet.layer3.1.conv2'      ,0  )
        g.addEdge('resnet.layer3.1.conv2'      ,'resnet.layer3.1.bn2'        ,0  )
        g.addEdge('resnet.layer3.1.bn2'        ,'resnet.layer3.1.relu2'      ,0  )
        g.addEdge('resnet.layer3.1.relu2'      ,'resnet.layer4.0.downsample' ,0  )
        g.addEdge('resnet.layer3.1.relu2'      ,'resnet.layer4.0.conv1'      ,0  )
        g.addEdge('resnet.layer4.0.downsample' ,'resnet.layer4.0.relu2'      ,0  )
        g.addEdge('resnet.layer4.0.conv1'      ,'resnet.layer4.0.bn1'        ,0  )
        g.addEdge('resnet.layer4.0.bn1'        ,'resnet.layer4.0.relu1'      ,0  )
        g.addEdge('resnet.layer4.0.relu1'      ,'resnet.layer4.0.conv2'      ,0  )
        g.addEdge('resnet.layer4.0.conv2'      ,'resnet.layer4.0.bn2'        ,0  )
        g.addEdge('resnet.layer4.0.bn2'        ,'resnet.layer4.0.relu2'      ,0  )
        g.addEdge('resnet.layer4.0.relu2'      ,'resnet.layer4.1.conv1'      ,0  )
        g.addEdge('resnet.layer4.0.relu2'      ,'resnet.layer4.1.relu2'      ,0  )
        g.addEdge('resnet.layer4.1.conv1'      ,'resnet.layer4.1.bn1'        ,0  )
        g.addEdge('resnet.layer4.1.bn1'        ,'resnet.layer4.1.relu1'      ,0  )
        g.addEdge('resnet.layer4.1.relu1'      ,'resnet.layer4.1.conv2'      ,0  )
        g.addEdge('resnet.layer4.1.conv2'      ,'resnet.layer4.1.bn2'        ,0  )
        g.addEdge('resnet.layer4.1.bn2'        ,'resnet.layer4.1.relu2'      ,0  )
        g.addEdge('resnet.layer4.1.relu2'      ,'resnet.avgpool'             ,0  )
        g.addEdge('resnet.avgpool'             ,'resnet.flatten'             ,0  )
        g.addEdge('resnet.flatten'             ,'resnet.fc'                  ,0  )
        
    def forward(self, x0, x1=0, start=0, end=58):
        Vertexs = list(self.g)
        cnt = 0
        for v in Vertexs:
            cnt += 1
            if cnt < start or cnt > end:
                continue
            # print(v.id, v.argc, len(v.getConnections()), cnt)
            if len(v.getConnections()) == 2:
                if v.argc == 1 or v.argc == 0:
                    x0 = v(x0)
                elif v.argc == 2:
                    x0 = v(x0, x1)
                x1 = x0 * 1
            elif 'downsample' in v.id:
                x1 = v(x0)
            else:
                if v.argc == 1 or v.argc == 0:
                    x0 = v(x0)
                elif v.argc == 2:
                    x0 = v(x0, x1)
        return x0, x1
    
    def get_cost(self, v: Vertex, A :list, B : list)->int:
        # print("****************")
        # print("this is node +" + str(v.id))
        
        cost = 0
        for i in range(len(B)):
            for j in list(B[i].keys()):
                # print("node is {}, cost is {}".format(j.id,B[i][j]))
                if(j.num > v.num):
                    weight = B[i][j]
                    cost = cost + weight        
        cost = cost + sum(A)
        return cost    
        
    def cut(self, n, dev):
    
        s = sum(vertex.cost for vertex in self.g)
        print(s)
             
        vertexs = list(self.g.getVertics())
        
        A = []
        B = []
        costs = []
        
        for i in range(0, len(vertexs)):
            A.append(self.g.getVertex(vertexs[i]).cost)
            
            if len(self.g.getVertex(vertexs[i]).connectedTo):
                B.append(self.g.getVertex(vertexs[i]).connectedTo)
                
            cost = self.get_cost(self.g.getVertex(vertexs[i]),A, B)
            costs.append(cost)
            
        costs.append((costs[-1] + 0))
        costs.insert(0,0)
        index = np.zeros(n, int)
        flag = 1
        for j in range(1, n ):
            c = dev[j - 1]/sum(dev)
            for i in range(index[j - 1], len(costs)-1):
                # print("cost[i] is {}, index[j] is {}".format(costs[i], index[j-1]))
                if costs[i] - costs[index[j-1]] > c*s:
                    if costs[ i - 1 ] - costs[index[j-1]] < c*s:
                        # index[j] = i - 1
                        if  flag:
                            index[j] = i -1
                        else:
                            index[j] = i             
                        flag = not flag
                        break
            
        
        index = np.append(index, len(vertexs))
        print(index)
        costt = []
        
        for i in range(0,len(index)-1):
            print("*********************************")
            print("this is block {}".format(str(i)))
            cost = 0
            for j in vertexs[index[i] : index[i + 1]]:
                cost = cost + self.g.getVertex(j).cost
                print(self.g.getVertex(j).id)
            costt.append(cost)
            
        # print(costt)
        

if __name__ =='__main__':
    g = Graph_resnet18(Timing=False)
    # warm up 
    for _ in range(5):
        x = torch.randn(1,3,224,224).to(device)
        g(x)
        
    g = Graph_resnet18(Timing=True)
    x = torch.randn(1,3,224,224).to(device)
    g(x)
    dev = [1/4.78, 1/3.60, 1/3.16, 1/2.60]
    g.cut(4, dev)
    
